import torch
import torchvision.transforms.functional as F


class ZeroPadToSize:
    def __init__(self, target_size=(210, 210)):
        self.target_size = target_size

    def __call__(self, image_tensor: torch.Tensor):

        _, height, width = image_tensor.size()
        padding_left = (self.target_size[1] - width) // 2
        padding_right = self.target_size[1] - width - padding_left
        padding_top = (self.target_size[0] - height) // 2
        padding_bottom = self.target_size[0] - height - padding_top

        return F.pad(image_tensor, padding=[padding_left, padding_top, padding_right, padding_bottom], fill=0)

class MinMaxNormalize:
    def __init__(self):
        pass

    def __call__(self, image_tensor: torch.Tensor):
        min_val = image_tensor.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
        max_val = image_tensor.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]

        normalized_image = (image_tensor - min_val) / (max_val - min_val)

        return normalized_image
    
class ManualCrop:
    def __init__(self, img_size):
        self.img_size = img_size

    def __call__(self, image):
        top = int(0.21 * self.img_size)
        left = int(0.325 * self.img_size)
        height = int(0.375 * self.img_size)
        width = int(0.375 * self.img_size)
        return F.crop(image, top=top, left=left, height=height, width=width)

class Resize:
    def __init__(self, target_size=(210, 210)):
        self.target_size = target_size

    def __call__(self, image):
        return F.resize(image, self.target_size)
    
class UnsqueezeDimension(object):
    def __init__(self, poisition:int) -> None:
        self.position = poisition
    def __call__(self, sample):
        return sample.unsqueeze_(self.position)
    
class SqueezeTensor(object):
    def __init__(self) -> None:
        pass
    def __call__(self, sample):
        return torch.squeeze_(sample)